#!/usr/bin/env python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import PIL
import time
import shlex
import shutil
import asyncio
import subprocess
import qemu.qmp
from ._util import Util
from ._const import Arch, Version


class VmConfig:

    def __init__(self, qemu_vm_type, cpu_arch, cpu_number, memory_size, bios_type, disk_interface, network_type):
        assert qemu_vm_type in ["pc", "q35"]
        assert cpu_arch in ["x86", "amd64"]
        assert cpu_number in [1, 2, 4]
        assert re.fullmatch("[0-9]+(M|G)", memory_size)
        assert bios_type in ["bios", "efi"]
        assert disk_interface in ["ide"]
        assert network_type in ["none", "isolated", "full"]

        assert not (cpu_arch == "x86" and bios_type == "efi")

        self.qemuVmType = qemu_vm_type                # str         "pc", "q35"
        self.cpuArch = cpu_arch                       # str
        self.cpuNumber = cpu_number                   # int
        self.memorySize = memory_size                 # str         "1M", "1G"
        self.biosType = bios_type                     # str
        self.diskInterface = disk_interface           # str         "ide", "virtio-blk"
        self.networkType = network_type               # str         "none", "isolated", "full"

    def export_to_file_as_ini_content(self):
        buf = "[main]\n"
        for k, v in self.__dict__.items():
            buf += "%s = %s\n" % (k, v)
        return buf


class Vm:

    def __init__(self, work_dir):
        self._workDirObj = work_dir

        from . import os_win98, os_winxp, os_win7                               # break circular dependencies, os_XXX module also uses Vm class
        if self._workDirObj._ts["version"] == Version.WINDOWS_98.value:
            self._vmCfg = os_win98.Spec().getVmConfig()
        elif self._workDirObj._ts["version"] == Version.WINDOWS_XP.value:
            self._vmCfg = os_winxp.Spec(Arch(self._workDirObj._ts["arch"])).getVmConfig()
        elif self._workDirObj._ts["version"] == Version.WINDOWS_7.value:
            self._vmCfg = os_win7.Spec(Arch(self._workDirObj._ts["arch"])).getVmConfig()
        else:
            assert False

        self._vncRecordFile = None

        self._noReboot = False

        self._cpuNumber = self._vmCfg.cpuNumber

        self._memSize = Util.sizeStrToInt(self._vmCfg.memorySize)
        assert self._memSize % (1024 * 1024) == 0

        self._biosFile = None

        self._cdromFileOrUrl = None
        self._cdromBoot = None

        self._floppyFile = None
        self._floppyBoot = None

        self._usbStickFileOrUrl = None

        self._networkType = self._vmCfg.networkType
        self._shareDir = None

    def get_cpu_number(self):
        return self._cpuNumber

    def get_memory_size(self):
        return self._memSize

    def get_network_servers(self):
        assert self._networkType in ["isolated", "full"]

        ret = {
            "dns": {
                "ipv4": "10.0.2.3",
                "ipv6": "xxx::3",
            },
        }

        if self._networkType == "full":
            ret["gateway"] = {
                "ipv4": "10.0.2.2",
                "ipv6": "xxx::2",
            }

        if self._shareDir is not None:
            ret["samba"] = {
                "ipv4": "10.0.2.4",
                "ipv6": "xxx::4",
            }

        return ret

    def set_record_file(self, path):
        assert not self.is_running()

        self._vncRecordFile = self._workDirObj._vmRecordFilePrefix + path

    def set_no_reboot(self, value):
        assert not self.is_running()

        self._noReboot = value

    def increase_cpu_number_to(self, number):
        assert not self.is_running()
        assert number > self._cpuNumber

        self._cpuNumber = number

    def enlarge_memory_to(self, size):
        assert not self.is_running()
        assert size % (1024 * 1024) == 0 and size > self._memSize

        self._memSize = size

    def use_custom_bios(self, bios_file):
        assert not self.is_running()

        self._biosFile = bios_file

    def use_floppy(self, floppy_index, floppy_file, boot=False):
        assert not self.is_running()
        assert floppy_index == 0

        self._floppyFile = floppy_file
        self._floppyBoot = boot

    def use_cdrom(self, filepath_or_url, boot=False):
        assert not self.is_running()

        self._cdromFileOrUrl = filepath_or_url
        self._cdromBoot = boot

    def use_usb_stick(self, filepath_or_url):
        assert not self.is_running()

        self._usbStickFileOrUrl = filepath_or_url

    def use_network(self, network_type):
        assert not self.is_running()

        # new network type must not be "none", must be higher than or equal to current network type
        if self._networkType == "none":
            assert network_type in ["isolated", "full"]
        elif self._networkType == "isolated":
            assert network_type in ["isolated", "full"]
        elif self._networkType == "full":
            assert network_type == "full"
        else:
            assert False

        self._networkType = network_type

    def share_host_dir(self, path):
        assert not self.is_running()
        assert self._networkType in ["isolated", "full"]
        assert path is not None

        self._shareDir = path

    def is_running(self):
        return hasattr(self, "_proc")

    def start(self, show=False):
        try:
            self._bShow = show
            self._cmdLine = self._generateQemuCommand()

            # save before start, so that we can inspect it on failure
            self._workDirObj._saveVm(self._vmCfg, self._cmdLine)

            # create efi var file if needed
            if self._vmCfg.biosType == "efi" and not os.path.exists(self._workDirObj._efiVarFile):
                shutil.copy("/usr/share/edk2/OvmfX64/OVMF_VARS.fd", self._workDirObj._efiVarFile)

            # start vm
            self._proc = subprocess.Popen(shlex.split(self._cmdLine.replace("\\\n", "")))
            self._waitVmStart()

            # start video recording
            # FIXME: ffmpeg does not support vnc, suck and strange
            # if self._vncRecordFile is not None:
            #    self._recordProc = self._createVideoRecordProcess()
        except BaseException:
            self.stop()
            raise

    def stop(self):
        if hasattr(self, "_recordProc"):
            self._recordProc.terminate()
            self._recordProc.wait()
            del self._recordProc
        if hasattr(self, "_proc"):
            self._proc.terminate()
            self._proc.wait()
            del self._proc
            del self._cmdLine
            del self._bShow

    def wait_until_stable(self):
        self._lastImg = None
        self._stableCount = None
        while self._stableCount is None or self._stableCount < 10:        # test if the system has stablized for 10 seconds
            if self._lastImg is None:
                self._lastImg = self.get_screen_content()
                self._stableCount = 0
            elif self._lastImg != self.get_screen_content():
                self._lastImg = None
                self._stableCount = None
            else:
                self._stableCount += 1
            time.sleep(1)

    def shutdown(self):
        async def _asyncFunc():
            c = qemu.qmp.QMPClient()
            await c.connect(self._workDirObj.vm_control_socket_filepath)
            await c.execute("system_powerdown")
            await c.disconnect()
        asyncio.run(_asyncFunc())

    def poweroff(self):
        self._proc.terminate()

    # deprecated
    def get_screen_resolution(self):
        return self.retrieve_screen_resolution()

    def retrieve_screen_resolution(self):
        # return (width, height)
        return self.retrieve_screen_content().size

    # deprecated
    def get_screen_content(self):
        return self.retrieve_screen_content()

    def retrieve_screen_content(self):
        self._dumpScreen(self._workDirObj._vmScreenDumpFile)
        return PIL.Image.open(self._workDirObj._vmScreenDumpFile)

    def input_keyboard(self, input_type, *kargs):
        if input_type == "keys":
            assert len(kargs) > 0
            keys = kargs
        elif input_type == "line":
            assert len(kargs) == 1
            keys = []
            for c in kargs[0]:
                if c == " ":
                    keys.append("spc")
                elif c == ".":
                    keys.append("dot")
                elif c == "-":
                    keys.append("minus")
                elif c == "\\":
                    keys.append("backslash")
                elif c == ":":
                    keys.append("shift+semicolon")
                elif re.fullmatch("[A-Z]", c):
                    keys.append("shift+" + c.lower())
                else:
                    keys.append(c)
            keys.append("kp_enter")
        else:
            assert False

        async def _asyncFunc():
            c = qemu.qmp.QMPClient()
            await c.connect(self._workDirObj.vm_control_socket_filepath)
            for key in keys:
                klist = [{"type": "qcode", "data": x} for x in key.split("+")]
                await c.execute("send-key", {
                    "keys": klist,
                })
                await asyncio.sleep(0.1)
            await c.disconnect()
        asyncio.run(_asyncFunc())

    def input_mouse(self, input_type, *kargs):
        if input_type == "abs-move":
            assert len(kargs) == 2
            x, y = kargs[0], kargs[1]
            btn = None
        elif input_type == "abs-move-and-click-left-button":
            assert len(kargs) == 2
            x, y = kargs[0], kargs[1]
            btn = "left"
        elif input_type == "abs-move-and-click-right-button":
            assert len(kargs) == 2
            x, y = kargs[0], kargs[1]
            btn = "right"
        else:
            assert False

        # it seems that usb-tablet has (x, y) range as (0, 32600) through trying, so we do a convert
        # but https://qemu.eu/doc/5.0/qemu-qmp-ref.html says is (0, 0x7ffff) ?
        w, h = self.get_screen_resolution()
        x = 32600 * x // w
        y = 32600 * y // h

        async def _asyncFunc():
            c = qemu.qmp.QMPClient()
            await c.connect(self._workDirObj.vm_control_socket_filepath)
            await c.execute("input-send-event", {
                "events": [
                    { "type": "abs", "data" : { "axis": "x", "value" : x } },
                    { "type": "abs", "data" : { "axis": "y", "value" : y } },
                ]
            })
            if btn is not None:
                # only take effect if sent in 2 seperate events
                await c.execute("input-send-event", {
                    "events": [
                        { "type": "btn", "data" : { "button": btn, "down": True }},
                    ],
                })
                await c.execute("input-send-event", {
                    "events": [
                        { "type": "btn", "data" : { "button": btn, "down": False }},
                    ],
                })

            await asyncio.sleep(0.1)
            await c.disconnect()
        asyncio.run(_asyncFunc())

    def change_floppy(self, floppy_index, floppy_file):
        assert False

    def change_cdrom(self, filepath_or_url):
        if filepath_or_url is not None:
            assert "://" not in filepath_or_url     # FIXME: should support url

            async def _asyncFunc():
                c = qemu.qmp.QMPClient()
                await c.connect(self._workDirObj.vm_control_socket_filepath)
                await c.execute("blockdev-change-medium", {
                    "device": "cdrom",
                    "filename": filepath_or_url,
                })
                await c.disconnect()
        else:
            async def _asyncFunc():
                c = qemu.qmp.QMPClient()
                await c.connect(self._workDirObj.vm_control_socket_filepath)
                await c.execute("eject", {
                    "device": "cdrom",
                })
                await c.disconnect()

        asyncio.run(_asyncFunc())

    def wait_until_stop(self):
        self._proc.wait()
        del self._proc
        del self._cmdLine
        del self._bShow

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, type, value, traceback):
        self.stop()

    def _generateQemuCommand(self):
        # qemu command
        if self._vmCfg.cpuArch == "x86":
            cmd = "qemu-system-i386 \\\n"
        elif self._vmCfg.cpuArch == "amd64":
            cmd = "qemu-system-x86_64 \\\n"
        else:
            assert False
        cmd += "    -no-user-config \\\n"
        cmd += "    -nodefaults \\\n"
        if self._noReboot:
            cmd += "    -no-reboot \\\n"

        # machine
        if True:
            optList = []
            optList.append(self._vmCfg.qemuVmType)
            optList.append("usb=on")
            if os.access("/dev/kvm", os.R_OK | os.W_OK):
                optList.append("accel=kvm")
            cmd += "    -machine %s \\\n" % (",".join(optList))

        # cpu
        cmd += "    -cpu host \\\n"
        cmd += "    -smp %d,sockets=1,cores=%d,threads=1 \\\n" % (self._cpuNumber, self._cpuNumber)

        # memory
        cmd += "    -m %dM \\\n" % (self._memSize // (1024 * 1024))

        # bios
        if self._vmCfg.biosType == "bios":
            if self._biosFile is not None:
                cmd += "    -bios %s \\\n" % (self._biosFile)
            else:
                pass
        elif self._vmCfg.biosType == "efi":
            assert self._biosFile is None
            cmd += "    -drive if=pflash,format=raw,readonly=on,file=/usr/share/edk2/OvmfX64/OVMF_CODE.fd \\\n"
            cmd += "    -drive if=pflash,format=raw,file=%s \\\n" % (self._workDirObj._efiVarFile)
        else:
            assert False

        # platform devices
        cmd += "    -rtc base=localtime \\\n"           # FIXME: how to do it in a more standardized way?

        # additional controllers
        if self._vmCfg.qemuVmType == "pc":
            pass
        elif self._vmCfg.qemuVmType == "q35":
            # FIXME: enable qemu-xhci causes virtual usb stick not recoginized, why?
            # cmd += "    -device qemu-xhci \\\n"                 # xhci controller supports USB1.1 ~ USB3.0, which is optimal, FIXME, does q35 already has it?
            cmd += "    -device isa-fdc \\\n"
        else:
            assert False

        # main-disk
        if True:
            # "cache.no-flush=on" is the fastest, we don't need safety
            cmd += "    -blockdev 'driver=file,filename=%s,node-name=main-disk,cache.no-flush=on' \\\n" % (self._workDirObj.image_filepath)
            if self._vmCfg.diskInterface == "ide":
                cmd += "    -device ide-hd,bus=ide.0,unit=0,drive=main-disk \\\n"
            elif self._vmCfg.diskInterface == "virtio-blk":
                cmd += "    -device virtio-blk,drive=main-disk \\\n"
            else:
                assert False

        # cdrom
        if self._cdromFileOrUrl is not None:
            # if self._cdromFileOrUrl.startswith("http://") or self._cdromFileOrUrl.startswith("https://"):
            #     sDriver = "driver=http,url=%s" % (self._cdromFileOrUrl)
            # elif self._cdromFileOrUrl.startswith("ftp://"):
            #     # FIXME: really?
            #     sDriver = "driver=http,url=%s" % (self._cdromFileOrUrl)
            # else:
            #     sDriver = "driver=file,filename=%s" % (self._cdromFileOrUrl)
            # cmd += "    -blockdev '%s,read-only=on,node-name=cdrom' \\\n" % (sDriver)
            # cmd += "    -device ide-cd,bus=ide.1,unit=0,drive=cdrom \\\n"
            cmd += "    -drive id=cdrom,if=ide,bus=1,unit=0,media=cdrom,readonly=on,file=%s \\\n" % (self._cdromFileOrUrl)      # FIXME: only -drive supports "eject" and "change medium", see https://gitlab.com/qemu-project/qemu/-/issues/678
        else:
            cmd += "    -drive id=cdrom,if=ide,bus=1,unit=0,media=cdrom,readonly=on \\\n"

        # floppy
        if self._floppyFile is not None:
            # use "unit=0" to make it undoubtly "A:"
            cmd += "    -blockdev 'driver=file,filename=%s,node-name=floppy' \\\n" % (self._floppyFile)
            cmd += "    -device floppy,bus=floppy-bus.0,unit=0,drive=floppy \\\n"
        else:
            # use "unit=0" to make it undoubtly "A:"
            cmd += "    -device floppy,bus=floppy-bus.0,unit=0 \\\n"

        if self._usbStickFileOrUrl is not None:
            cmd += "    -blockdev 'driver=file,filename=%s,node-name=usbstick' \\\n" % (self._usbStickFileOrUrl)
            cmd += "    -device usb-storage,drive=usbstick,removable=on \\\n"

        # graphics device
        if self._bShow:
            cmd += "    -display sdl \\\n"
        cmd += "    -device VGA \\\n"
        cmd += "    -vnc unix:%s \\\n" % (self._workDirObj.vm_display_socket_filepath)

        # use tablet instead of mouse, so that we can use absolute coordination to do automated operations
        if True:
            cmd += "    -device usb-tablet \\\n"

        # network device
        if self._networkType in ["full", "isolated"]:
            subcmd = ""
            if self._networkType == "isolated":
                subcmd += ",restrict=on"
            if self._shareDir is not None:
                subcmd += ",smb=%s" % (os.path.abspath(self._shareDir))
            cmd += "    -netdev user,id=eth0%s \\\n" % (subcmd)
            cmd += "    -device rtl8139,netdev=eth0,romfile= \\\n"
        elif self._networkType == "none":
            pass
        else:
            assert False

        # monitor interface
        if True:
            cmd += "    -chardev socket,id=qmp,server=on,wait=off,path=%s \\\n" % (self._workDirObj.vm_control_socket_filepath)
            cmd += "    -mon chardev=qmp,mode=control \\\n"

        # boot order
        if self._cdromBoot:
            cmd += "    -boot order=dc \\\n"
        elif self._floppyBoot:
            cmd += "    -boot order=ac \\\n"
        else:
            cmd += "    -boot order=c \\\n"

        # eliminate the last " \\\n"
        cmd = cmd[:-3] + "\n"

        return cmd

    def _waitVmStart(self):
        while True:
            time.sleep(0.1)
            if not os.path.exists(self._workDirObj.vm_display_socket_filepath):
                continue
            if not os.path.exists(self._workDirObj.vm_control_socket_filepath):
                continue
            break

    def _createVideoRecordProcess(self):
        cmd = [
            "ffmpeg",
            "-i", "unix://" + self._workDirObj.vm_display_socket_filepath,
            self._vncRecordFile,
        ]
        return subprocess.Popen(cmd)

    def _dumpScreen(self, dstFile):
        # QEMU SUCKS: qmp command screendump returns before dump finishes!!!
        # so we have to use ugly code to detect whether the dump is finished

        tmpDstFile = dstFile + ".tmp"

        async def _asyncFunc():
            c = qemu.qmp.QMPClient()
            await c.connect(self._workDirObj.vm_control_socket_filepath)
            await c.execute("screendump", {"filename": tmpDstFile})
            await c.disconnect()
        asyncio.run(_asyncFunc())

        def _has_process_using_file(file_path):
            try:
                subprocess.check_output(["fuser", file_path])
                return True                                     # return code 0 means there is processes using this file
            except subprocess.CalledProcessError:
                return False                                    # return code 1 means there is no process using this file
        while True:
            if not os.path.exists(tmpDstFile) or _has_process_using_file(tmpDstFile):
                print("debug sucks 1")
                time.sleep(0.1)
                continue
            os.rename(tmpDstFile, dstFile)
            break
